import numpy as np
import tensorflow as tf
from tensorflow.keras.models import Sequential
from tensorflow.keras.layers import Conv2D, MaxPooling2D, Flatten, Dense, Dropout
from tensorflow.keras.preprocessing.image import ImageDataGenerator
# 准备数据集
train_data_dir = 'D:/datas/training_data'  # 训练数据文件夹路径
validation_data_dir = 'D:/datas/validation_data'  # 验证数据文件夹路径
img_width, img_height = 500, 400  # 图像尺寸
batch_size = 32
# 数据增强和预处理
train_datagen = ImageDataGenerator(
    rescale=1./255,
    shear_range=0.2,
    zoom_range=0.2,
    horizontal_flip=True)
test_datagen = ImageDataGenerator(rescale=1./255)
train_generator = train_datagen.flow_from_directory(
    train_data_dir,
    target_size=(img_width, img_height),
    batch_size=batch_size,
    class_mode='categorical')
validation_generator = test_datagen.flow_from_directory(
    validation_data_dir,
    target_size=(img_width, img_height),
    batch_size=batch_size,
    class_mode='categorical')
# 定义模型
def create_model(input_shape, num_classes):
    model = Sequential([
        Conv2D(32, (3, 3), activation='relu', input_shape=input_shape),
        MaxPooling2D((2, 2)),
        Conv2D(64, (3, 3), activation='relu'),
        MaxPooling2D((2, 2)),
        Conv2D(128, (3, 3), activation='relu'),
        MaxPooling2D((2, 2)),
        Flatten(),
        Dense(128, activation='relu'),
        Dropout(0.5),
        Dense(num_classes, activation='softmax')
    ])
    return model
# 模型参数
input_shape = (img_width, img_height, 3)
num_classes = 5  # 有机垃圾、可回收垃圾、有害垃圾和其他垃圾四个类别
# 创建模型
model = create_model(input_shape, num_classes)
# 编译模型
model.compile(optimizer='adam',
              loss='categorical_crossentropy',
              metrics=['accuracy'])
# 训练模型
epochs = 10
history = model.fit(
    train_generator,
    steps_per_epoch=train_generator.samples // batch_size,
    epochs=epochs,
    validation_data=validation_generator,
    validation_steps=validation_generator.samples // batch_size)
# 保存模型
model.save('garbage_classifier_model.h5')
# 标签映射
label_map = {
    0: '硬纸板',
    1: '玻璃',
    2: '金属',
    3: '纸张',
    4: '塑料',
}
# 使用模型进行预测
def predict_image_class(image_path):
    img = tf.keras.preprocessing.image.load_img(image_path, target_size=(img_width, img_height))
    img_array = tf.keras.preprocessing.image.img_to_array(img)
    img_array = np.expand_dims(img_array, axis=0)  # 添加一个维度作为 batch
    predictions = model.predict(img_array)
    predicted_class_index = np.argmax(predictions, axis=1)
    predicted_class_label = label_map[predicted_class_index[0]]
    return predicted_class_label
# 示例用法
image_path = 'D:/datas/test_image.jpg'
predicted_class = predict_image_class(image_path)
print('Predicted class:', predicted_class)